# -*- coding: utf-8 -*-
"""
数据加载工具模块

提供统一的数据加载接口，用于加载问题数据和其他测试数据。
"""

import json
import os
from typing import List
import random
from loguru import logger


class DataLoader:
    """
    简单的数据加载器，用于加载问题数据
    """
    
    def __init__(self, data_dir: str = None):
        """
        初始化数据加载器
        
        Args:
            data_dir: 数据目录路径，默认为当前文件所在目录
        """
        if data_dir is None:
            # 默认使用当前文件所在目录
            self.data_dir = os.path.dirname(__file__)
        else:
            self.data_dir = data_dir
    
    def load_questions(self, filename: str = "questions.json") -> List[str]:
        """
        从JSON文件加载问题列表
        
        Args:
            filename: 问题文件名，默认为 "questions.json"
            
        Returns:
            问题列表
        """
        file_path = os.path.join(self.data_dir, filename)
        
        # 检查文件是否存在
        if not os.path.exists(file_path):
            logger.warning(f"问题文件不存在: {file_path}")
            return []
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 简单处理不同的数据格式
            questions = []
            if isinstance(data, list):
                questions = data
            elif isinstance(data, dict):
                # 尝试常见的字段名
                for key in ['questions', 'data', 'items']:
                    if key in data:
                        questions = data[key]
                        break
                else:
                    # 如果没有找到，返回所有值
                    questions = list(data.values())
            else:
                logger.warning(f"不支持的数据格式: {type(data)}")
                return []
            
            # 过滤掉空字符串和None值
            filtered_questions = []
            for q in questions:
                if q is not None and str(q).strip():
                    filtered_questions.append(str(q).strip())
            
            return filtered_questions
                
        except Exception as e:
            logger.warning(f"加载文件失败: {e}")
            return []
    
    def load_sample_questions(self, filename: str = "questions.json", sample_size: int = 10) -> List[str]:
        """
        加载样本问题
        
        Args:
            filename: 问题文件名
            sample_size: 样本大小
            
        Returns:
            样本问题列表
        """
        all_questions = self.load_questions(filename)
        
        if len(all_questions) <= sample_size:
            return all_questions
        
        return random.sample(all_questions, sample_size)
    
    def get_test_queries(self, questions: List[str], count: int = 5) -> List[str]:
        """
        从问题列表中随机选择测试查询
        
        Args:
            questions: 问题列表
            count: 选择的查询数量
            
        Returns:
            测试查询列表
        """
        if len(questions) <= count:
            return questions.copy()
        
        return random.sample(questions, count)


# 便捷函数
def load_questions(filename: str = "questions.json", data_dir: str = None) -> List[str]:
    """
    便捷函数：直接加载问题数据
    
    Args:
        filename: 问题文件名
        data_dir: 数据目录，默认为当前模块目录
        
    Returns:
        问题列表
    """
    loader = DataLoader(data_dir)
    return loader.load_questions(filename)

def load_sample_questions(sample_size: int = 10, filename: str = "questions.json", data_dir: str = None) -> List[str]:
    """
    便捷函数：加载样本问题
    
    Args:
        sample_size: 样本大小
        filename: 问题文件名
        data_dir: 数据目录
        
    Returns:
        样本问题列表
    """
    loader = DataLoader(data_dir)
    return loader.load_sample_questions(filename, sample_size)

def get_test_queries(questions: List[str], count: int = 5) -> List[str]:
    """
    便捷函数：获取测试查询
    
    Args:
        questions: 问题列表
        count: 查询数量
        
    Returns:
        测试查询列表
    """
    loader = DataLoader()
    return loader.get_test_queries(questions, count)


if __name__ == "__main__":
    # 测试数据加载
    try:
        questions = load_questions()
        logger.info(f"加载了 {len(questions)} 个问题")
        logger.info("前5个问题:")
        for i, q in enumerate(questions[:5], 1):
            logger.info(f"{i}. {q}")
        
        # 测试示例问题加载
        sample_questions = load_sample_questions(5)
        logger.info(f"\n示例问题 ({len(sample_questions)} 个):")
        for i, q in enumerate(sample_questions, 1):
            logger.info(f"{i}. {q}")
        
        # 测试查询生成
        test_queries = get_test_queries(questions, 3)
        logger.info(f"\n测试查询 ({len(test_queries)} 个):")
        for i, q in enumerate(test_queries, 1):
            logger.info(f"{i}. {q}")
            
    except Exception as e:
        logger.error(f"测试失败: {e}")